"""
Task Queue System for Multi-User Gradio

GPU
"""
import queue
import threading
import time
import uuid
import logging
from typing import Callable, Any, Dict, Optional
from dataclasses import dataclass
from datetime import datetime
from enum import Enum

logger = logging.getLogger(__name__)


class TaskStatus(Enum):
    """"""
    PENDING = "pending"      # 
    PROCESSING = "processing"  # 
    COMPLETED = "completed"   # 
    FAILED = "failed"         # 
    CANCELLED = "cancelled"   # 


@dataclass
class Task:
    """"""
    task_id: str
    func: Callable
    args: tuple
    kwargs: dict
    status: TaskStatus
    result: Any = None
    error: Optional[str] = None
    created_at: float = None
    started_at: Optional[float] = None
    completed_at: Optional[float] = None
    position: int = 0  # 

    def __post_init__(self):
        if self.created_at is None:
            self.created_at = time.time()

    @property
    def wait_time(self) -> float:
        """"""
        if self.started_at:
            return self.started_at - self.created_at
        return time.time() - self.created_at

    @property
    def processing_time(self) -> Optional[float]:
        """"""
        if self.completed_at and self.started_at:
            return self.completed_at - self.started_at
        return None

    @property
    def total_time(self) -> Optional[float]:
        """"""
        if self.completed_at:
            return self.completed_at - self.created_at
        return None


class TaskQueue:
    """
    

    
    - FIFO
    - workerGPU
    - 
    - 
    """

    def __init__(self, max_history: int = 100):
        """
        Args:
            max_history: 
        """
        self.task_queue = queue.Queue()
        self.tasks: Dict[str, Task] = {}  # task_id -> Task
        self.max_history = max_history
        self.worker_thread = None
        self.running = False
        self.lock = threading.Lock()

        # 
        self.stats = {
            'total_tasks': 0,
            'completed_tasks': 0,
            'failed_tasks': 0,
            'total_wait_time': 0.0,
            'total_processing_time': 0.0,
        }

    def start(self):
        """worker"""
        if self.running:
            logger.warning("Task queue already running")
            return

        self.running = True
        self.worker_thread = threading.Thread(target=self._worker, daemon=True)
        self.worker_thread.start()
        logger.info(" Task queue started")

    def stop(self):
        """worker"""
        self.running = False
        if self.worker_thread:
            self.worker_thread.join(timeout=5)
        logger.info(" Task queue stopped")

    def submit(self, func: Callable, *args, **kwargs) -> str:
        """
        

        Args:
            func: 
            *args, **kwargs: 

        Returns:
            task_id: ID
        """
        task_id = str(uuid.uuid4())[:8]  # ID

        task = Task(
            task_id=task_id,
            func=func,
            args=args,
            kwargs=kwargs,
            status=TaskStatus.PENDING
        )

        with self.lock:
            self.tasks[task_id] = task
            self.stats['total_tasks'] += 1

        # 
        self.task_queue.put(task_id)

        # 
        self._update_positions()

        logger.info(f" Task {task_id} submitted, queue size: {self.get_queue_size()}")

        return task_id

    def _worker(self):
        """Worker"""
        logger.info(" Worker thread started")

        while self.running:
            try:
                # 1
                task_id = self.task_queue.get(timeout=1.0)

                with self.lock:
                    task = self.tasks.get(task_id)

                if task is None:
                    logger.warning(f"Task {task_id} not found")
                    continue

                # 
                self._execute_task(task)

                # 
                self._update_positions()

            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"Worker error: {e}", exc_info=True)

        logger.info(" Worker thread stopped")

    def _execute_task(self, task: Task):
        """"""
        task_id = task.task_id

        try:
            # 
            task.status = TaskStatus.PROCESSING
            task.started_at = time.time()

            logger.info(f"  Processing task {task_id}")

            # 
            result = task.func(*task.args, **task.kwargs)

            # 
            task.result = result
            task.status = TaskStatus.COMPLETED
            task.completed_at = time.time()

            # 
            with self.lock:
                self.stats['completed_tasks'] += 1
                self.stats['total_wait_time'] += task.wait_time
                if task.processing_time:
                    self.stats['total_processing_time'] += task.processing_time

            logger.info(
                f" Task {task_id} completed in {task.processing_time:.2f}s "
                f"(waited {task.wait_time:.2f}s)"
            )

        except Exception as e:
            # 
            task.status = TaskStatus.FAILED
            task.error = str(e)
            task.completed_at = time.time()

            with self.lock:
                self.stats['failed_tasks'] += 1

            logger.error(f" Task {task_id} failed: {e}", exc_info=True)

        finally:
            # 
            self._cleanup_old_tasks()

    def _update_positions(self):
        """"""
        pending_tasks = [
            (task_id, task) for task_id, task in self.tasks.items()
            if task.status == TaskStatus.PENDING
        ]

        # 
        pending_tasks.sort(key=lambda x: x[1].created_at)

        for position, (task_id, task) in enumerate(pending_tasks, start=1):
            task.position = position

    def _cleanup_old_tasks(self):
        """"""
        with self.lock:
            completed = [
                (task_id, task) for task_id, task in self.tasks.items()
                if task.status in (TaskStatus.COMPLETED, TaskStatus.FAILED)
            ]

            if len(completed) > self.max_history:
                # 
                completed.sort(key=lambda x: x[1].completed_at or 0)

                # 
                to_delete = len(completed) - self.max_history
                for task_id, _ in completed[:to_delete]:
                    del self.tasks[task_id]

    def get_task_status(self, task_id: str) -> Optional[Dict]:
        """"""
        task = self.tasks.get(task_id)
        if task is None:
            return None

        return {
            'task_id': task.task_id,
            'status': task.status.value,
            'position': task.position if task.status == TaskStatus.PENDING else 0,
            'wait_time': task.wait_time,
            'processing_time': task.processing_time,
            'total_time': task.total_time,
            'error': task.error,
        }

    def wait_for_result(self, task_id: str, timeout: float = 300) -> Any:
        """
        

        Args:
            task_id: ID
            timeout: 

        Returns:
            

        Raises:
            TimeoutError: 
            RuntimeError: 
        """
        start_time = time.time()

        while True:
            task = self.tasks.get(task_id)

            if task is None:
                raise RuntimeError(f"Task {task_id} not found")

            if task.status == TaskStatus.COMPLETED:
                return task.result

            if task.status == TaskStatus.FAILED:
                raise RuntimeError(f"Task failed: {task.error}")

            if time.time() - start_time > timeout:
                raise TimeoutError(f"Task {task_id} timeout after {timeout}s")

            # 
            time.sleep(0.5)

    def get_queue_size(self) -> int:
        """"""
        return self.task_queue.qsize()

    def get_stats(self) -> Dict:
        """"""
        with self.lock:
            stats = self.stats.copy()

        # 
        if stats['completed_tasks'] > 0:
            stats['avg_wait_time'] = stats['total_wait_time'] / stats['completed_tasks']
            stats['avg_processing_time'] = stats['total_processing_time'] / stats['completed_tasks']
        else:
            stats['avg_wait_time'] = 0.0
            stats['avg_processing_time'] = 0.0

        stats['queue_size'] = self.get_queue_size()
        stats['active_tasks'] = len([
            t for t in self.tasks.values()
            if t.status in (TaskStatus.PENDING, TaskStatus.PROCESSING)
        ])

        return stats

    def get_all_tasks(self) -> list:
        """"""
        return [
            {
                'task_id': task.task_id,
                'status': task.status.value,
                'created_at': datetime.fromtimestamp(task.created_at).isoformat(),
                'position': task.position,
            }
            for task in sorted(self.tasks.values(), key=lambda t: t.created_at, reverse=True)
        ]


# 
_task_queue = None

def get_task_queue() -> TaskQueue:
    """"""
    global _task_queue
    if _task_queue is None:
        _task_queue = TaskQueue(max_history=100)
        _task_queue.start()
    return _task_queue


if __name__ == '__main__':
    # 
    import random

    logging.basicConfig(level=logging.INFO)

    def dummy_task(task_num: int, sleep_time: float):
        """"""
        print(f"  Task {task_num} processing...")
        time.sleep(sleep_time)
        return f"Task {task_num} result"

    print("=" * 80)
    print("Task Queue Test")
    print("=" * 80)

    queue = get_task_queue()

    # 
    task_ids = []
    for i in range(5):
        sleep_time = random.uniform(1, 3)
        task_id = queue.submit(dummy_task, i, sleep_time)
        task_ids.append(task_id)
        print(f"Submitted task {i} (ID: {task_id})")

    print(f"\nQueue size: {queue.get_queue_size()}")
    print("\nWaiting for all tasks to complete...\n")

    # 
    for task_id in task_ids:
        try:
            result = queue.wait_for_result(task_id, timeout=30)
            print(f" {task_id}: {result}")
        except Exception as e:
            print(f" {task_id}: {e}")

    print("\n" + "=" * 80)
    print("Statistics:")
    print("=" * 80)
    stats = queue.get_stats()
    for key, value in stats.items():
        print(f"  {key}: {value}")

    queue.stop()
